// Physical memory allocator, for user processes,
// kernel stacks, page-table pages,
// and pipe buffers. Allocates whole 4096-byte pages.

#include "types.h"
#include "param.h"
#include "memlayout.h"
#include "spinlock.h"
#include "riscv.h"
#include "defs.h"

void freerange(void *pa_start, void *pa_end);

extern char end[]; // first address after kernel.
                   // defined by kernel.ld.

struct run {
  struct run *next;
};

struct {
  struct spinlock lock;
  struct run *freelist;
} kmems[NCPU];

void
kinit()
{
  char name[16];
  for (int i = 0; i < NCPU; i++) {
    snprintf(name, sizeof(name), "kmem%d", i);
    initlock(&kmems[i].lock, name);
  }
  freerange(end, (void*)PHYSTOP);
}

void
freerange(void *pa_start, void *pa_end)
{
  char *p;
  p = (char*)PGROUNDUP((uint64)pa_start);
  for(; p + PGSIZE <= (char*)pa_end; p += PGSIZE)
    kfree(p);
}

// Free the page of physical memory pointed at by pa,
// which normally should have been returned by a
// call to kalloc().  (The exception is when
// initializing the allocator; see kinit above.)
void
kfree(void *pa)
{
  struct run *r;

  if(((uint64)pa % PGSIZE) != 0 || (char*)pa < end || (uint64)pa >= PHYSTOP)
    panic("kfree");

  // 填充释放的页面以捕获错误
  memset(pa, 1, PGSIZE);

  r = (struct run*)pa;

  push_off(); // 关闭中断以安全调用 cpuid()
  int cpu = cpuid();
  acquire(&kmems[cpu].lock);
  r->next = kmems[cpu].freelist;
  kmems[cpu].freelist = r;
  release(&kmems[cpu].lock);
  pop_off(); // 恢复中断
}

// Allocate one 4096-byte page of physical memory.
// Returns a pointer that the kernel can use.
// Returns 0 if the memory cannot be allocated.
void *
kalloc(void)
{
  struct run *r;

  push_off();
  int cpu = cpuid();
  acquire(&kmems[cpu].lock);
  r = kmems[cpu].freelist;
  if (r) {
    kmems[cpu].freelist = r->next;
    release(&kmems[cpu].lock);
    pop_off();
    memset((char *)r, 5, PGSIZE);
    return (void *)r;
  }
  release(&kmems[cpu].lock); // 释放当前 CPU 锁，避免嵌套

  // 尝试窃取
  int steal_left = 16; // 减少窃取数量，降低竞争
  for (int i = 0; i < NCPU; i++) {
    if (i == cpu)
      continue;
    acquire(&kmems[i].lock);
    struct run *rr = kmems[i].freelist;
    struct run *prev = 0;
    int count = 0;
    // 计算可窃取的页面数
    while (rr && count < steal_left) {
      prev = rr;
      rr = rr->next;
      count++;
    }
    if (prev) {
      // 窃取链表的前半部分
      r = kmems[i].freelist;
      kmems[i].freelist = rr;
      prev->next = 0;
      release(&kmems[i].lock);
      acquire(&kmems[cpu].lock);
      prev->next = kmems[cpu].freelist;
      kmems[cpu].freelist = r;
      r = kmems[cpu].freelist;
      kmems[cpu].freelist = r->next;
      release(&kmems[cpu].lock);
      pop_off();
      memset((char *)r, 5, PGSIZE);
      return (void *)r;
    }
    release(&kmems[i].lock);
  }

  pop_off();
  return 0; // 无可用内存
}